/***************************************************************************
 *   Copyright (C) 2010 by Oleg Goncharov  *
 *   $EMAIL$                           *                          
 *                                                                         *
 *   This file is part of ChessVision.                                     *
 *                                                                         *
 *   ChessVision is free software; you can redistribute it and/or modify   *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/
#include "cpointpyrmatch.h"

#include <opencv/cv.h>
//#include <iostream>

using namespace std;

CPointPyrMatch::CPointPyrMatch(const Mat& sample, const Mat& test_img, double threshold) 
{
    d_maxlevel = floor(log2(MAX(sample.cols, sample.rows))) - 2;
    if (d_maxlevel < 0) 
    {
	d_maxlevel = 0;
    }

    d_sample = new Mat[d_maxlevel + 1];
    d_sample[0] = sample.clone();
    BuildPyr(d_sample);

    d_threshold = new double[d_maxlevel + 1];
    d_img = new Mat[d_maxlevel + 1];

    WCShiftAnalysis(test_img);
    for (int i = 0; i <= d_maxlevel; i++) 
    {
	d_threshold[i] *= threshold;
    }
}

/**
 * @fn void WCShiftAnalysis(const Mat& test_img)
 * Perform worst-case analysis pyramid thresholds tuning. It writes results to d_threshold and may decrease d_maxlevel.
 */
void CPointPyrMatch::WCShiftAnalysis(const Mat& test_img) 
{
    Mat corr;
    Size shifted_size(test_img.cols -(1<<d_maxlevel) + 1, test_img.rows - (1<<d_maxlevel) + 1);
    double min, max;

    for (int level = 0; level <= d_maxlevel; level++) 
    {
	d_threshold[level] = 1.0;
    }

    for (int dx = (1<<d_maxlevel) - 1; dx >=0 ; dx--)
    {
        for (int dy = (1<<d_maxlevel) - 1; dy >=0 ; dy--) 
	{
            d_img[0] = Mat(test_img, Rect(dx, dy, shifted_size.width, shifted_size.height));
            BuildPyr(d_img);

            for (int level = 1; level <= d_maxlevel; level++) 
	    {
                if (dx < 1<<level && dy < 1<<level) 
		{
                    matchTemplate(d_img[level], d_sample[level], corr, CV_TM_CCORR_NORMED);
                    minMaxLoc(corr, &min, &max);
                    if (max < 0.2) 
		    {
			d_maxlevel = level - 1;
		    }
                    if (max < d_threshold[level]) 
		    {
			d_threshold[level] = max;
		    }
                }
            }
        }
    }
}

CPointPyrMatch::~CPointPyrMatch() 
{
    delete [] d_sample;
    delete [] d_img;
    delete [] d_threshold;
}

bool CPointPyrMatch::DetectPoint(const Mat& img, Point2f& point)
{
    d_img[0] = img;
    BuildPyr(d_img);
    //Mat corr(d_img[d_maxlevel].size()-d_sample[d_maxlevel].size()+Size(1,1), CV_32FC1);
    Mat corr;
    Point min_loc, max_loc, best_max_loc;
    Size sample_size = d_sample[d_maxlevel].size();
    double min, max, best_max = 0.0;

    matchTemplate(d_img[d_maxlevel], d_sample[d_maxlevel], corr, CV_TM_CCORR_NORMED);
    
    for (int attempt = 0;  attempt < 8; attempt++) 
    {
        do 
	{
            minMaxLoc(corr, &min, &max, &min_loc, &max_loc);
            if (max < d_threshold[d_maxlevel]) 
	    {
                if (best_max < d_threshold[0]) return false;

                point.x = best_max_loc.x + (d_sample[0].cols>>1);
                point.y = best_max_loc.y + (d_sample[0].rows>>1);
                return true;
            }

       //     std::cout << "maxlevel max " << max << max_loc << std::endl;

            Rect rect(max_loc, sample_size);

            max = RefineMatch(max_loc);

            rect -= Point(sample_size.width>>1, sample_size.height>>1);
            rect &= Rect(0, 0, corr.cols, corr.rows);
       //     std::cout << "[" << rect.x << " " << rect.y << " " << rect.width << " " << rect.height << std::endl;
            corr(rect) = 0.0;
        } while (max < d_threshold[0]);
      //  std::cout << "0 max " << max << max_loc << std::endl;
        if (best_max < max) 
	{
            best_max = max;
            best_max_loc = max_loc;
        }
    }

//     std::cout << "best max " << best_max << std::endl;

    point.x = best_max_loc.x + d_sample[0].cols/2;
    point.y = best_max_loc.y + d_sample[0].rows/2;
    return true;
}

void CPointPyrMatch::BuildPyr(Mat * pyr)
{
    for (int i = 1; i <= d_maxlevel; i++)     
    {
	pyrDown(pyr[i-1], pyr[i]);
    }
}

double CPointPyrMatch::RefineMatch(Point& point) {
    int d_border = 2;
    Size prox_size(2*d_border + 1, 2*d_border + 1);
    Size delta_size;
    Rect subimg_rect;
    Point min_loc, max_loc;
    Mat corr;
    double min, max;

    for (int level = d_maxlevel - 1; level >= 0; level--) 
    {
        point.x *= 2;
        point.y *= 2;

        delta_size = d_sample[level].size() - Size(1,1);

        subimg_rect = Rect(point, prox_size + delta_size);
        subimg_rect -= Point(d_border, d_border);
        subimg_rect &= Rect(Point(0,0), d_img[level].size());

        matchTemplate(d_img[level](subimg_rect), d_sample[level], corr,
                          CV_TM_CCORR_NORMED);

    //imshow("img", d_img[level](subimg_rect));
    //imshow("compare", corr);
    //imshow("sample", d_sample[level]);
    //waitKey(0);
    
        minMaxLoc(corr, &min, &max, &min_loc, &max_loc);
        if (max < d_threshold[level]) return 0.0;

        point = max_loc + subimg_rect.tl();
    }
    return max;
}

bool CPointPyrMatch::GetSampleSize(Size& size)
{
    if (d_sample)
    {	    
	size = d_sample[0].size();
	return 0;
    }	
    return 1;	
};